{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "## Imports and initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "import io\n",
    "import itertools\n",
    "\n",
    "import numpy as np\n",
    "import PIL.Image\n",
    "import requests\n",
    "import rerun as rr  # pip install rerun-sdk\n",
    "import rerun.blueprint as rrb\n",
    "import torch\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "## Define neural field class\n",
    "\n",
    "First, we define the neural field class which we can be used to represent any continuous ND signal. I.e., it maps an ND point to another ND point. In this notebook we fit fields to map from 2D image coordinates to RGB colors. This way the network weights can be interpreted as encoding a continuous image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "class NeuralField(torch.nn.Module):  # type: ignore[misc]\n",
    "    \"\"\"Simple neural field composed of positional encoding, MLP, and activation function.\"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        num_layers: int,\n",
    "        dim_hidden: int,\n",
    "        dim_in: int = 2,\n",
    "        dim_out: int = 3,\n",
    "        activation: str = \"sigmoid\",\n",
    "        pe_sigma: float | None = None,\n",
    "    ) -> None:\n",
    "        super().__init__()\n",
    "\n",
    "        self.num_layers = num_layers\n",
    "        self.dim_hidden = dim_hidden\n",
    "        self.dim_in = dim_in\n",
    "        self.dim_out = dim_out\n",
    "        self.activation = activation\n",
    "        self.pe_sigma = pe_sigma\n",
    "\n",
    "        sizes = [dim_in] + [dim_hidden for _ in range(num_layers - 1)] + [dim_out]\n",
    "        self.linears = torch.nn.ModuleList()\n",
    "        for in_size, out_size in itertools.pairwise(sizes):\n",
    "            self.linears.append(torch.nn.Linear(in_size, out_size))\n",
    "\n",
    "        if self.pe_sigma is not None:\n",
    "            torch.nn.init.normal_(self.linears[0].weight, 0.0, self.pe_sigma)\n",
    "\n",
    "    def __str__(self) -> str:\n",
    "        return f\"{self.num_layers} lay., {self.dim_hidden} neu., pe σ: {self.pe_sigma}\"\n",
    "\n",
    "    def forward(self, input_points: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Compute output for given input points.\n",
    "\n",
    "        Args:\n",
    "            input_points: input points\n",
    "\n",
    "        \"\"\"\n",
    "        if self.pe_sigma is None:\n",
    "            out = torch.relu(self.linears[0](input_points))\n",
    "        else:\n",
    "            out = torch.sin(self.linears[0](input_points))\n",
    "\n",
    "        for linear in self.linears[1:-1]:\n",
    "            out = torch.relu(linear(out))\n",
    "\n",
    "        out = self.linears[-1](out)\n",
    "\n",
    "        if self.activation == \"sigmoid\":\n",
    "            out = torch.sigmoid(out)\n",
    "\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "## Initialize and visualize neural fields\n",
    "\n",
    "Now we create a few neural fields with different parameters and visualize their output as images. We assume that images are fit in a 0 to 1 unit square, so we query in a dense grid (with some additional margin to observe out-of-training behavior) to retrieve the image from the network. Note that the positional encoding encodes how quickly the neural field varies out-of-the-box. This corresponds to the amount of detail that the field can easily represent, but also determines how the field extrapolates outside of the training region."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "fields = [\n",
    "    NeuralField(num_layers=5, dim_hidden=128, pe_sigma=5),\n",
    "    NeuralField(num_layers=5, dim_hidden=128, pe_sigma=15),\n",
    "    NeuralField(num_layers=5, dim_hidden=128, pe_sigma=30),\n",
    "    NeuralField(num_layers=5, dim_hidden=128, pe_sigma=100),\n",
    "]\n",
    "total_iterations = [0 for _ in fields]\n",
    "\n",
    "rr.init(\"rerun_example_cube\")\n",
    "\n",
    "blueprint = rrb.Blueprint(\n",
    "    rrb.Vertical(\n",
    "        rrb.Grid(\n",
    "            rrb.Spatial2DView(name=\"Target\", origin=\"target\"),\n",
    "            *[rrb.Spatial2DView(name=str(field), origin=f\"field_{i}\") for i, field in enumerate(fields)],\n",
    "        ),\n",
    "        rrb.TimeSeriesView(\n",
    "            name=\"Losses\",\n",
    "            origin=\"/\",\n",
    "            plot_legend=rrb.Corner2D.LeftTop,\n",
    "        ),\n",
    "        row_shares=[0.7, 0.3],\n",
    "    ),\n",
    "    collapse_panels=True,\n",
    ")\n",
    "for i, field in enumerate(fields):\n",
    "    rr.log(\n",
    "        f\"loss/field_{i}\",\n",
    "        rr.SeriesLines(names=str(field), aggregation_policy=rr.components.AggregationPolicy.auto(\"Average\")),\n",
    "        static=True,\n",
    "    )\n",
    "\n",
    "rr.notebook_show(blueprint=blueprint, width=1050, height=600)\n",
    "\n",
    "\n",
    "@torch.no_grad()  # type: ignore[misc]\n",
    "def log_field_as_image(\n",
    "    entity_path: str,\n",
    "    field: NeuralField,\n",
    "    min_uv: tuple[float, float],\n",
    "    max_uv: tuple[float, float],\n",
    "    uv_resolution: tuple[int, int],\n",
    ") -> None:\n",
    "    u_values = torch.linspace(min_uv[0], max_uv[0], uv_resolution[0])\n",
    "    v_values = torch.linspace(min_uv[1], max_uv[1], uv_resolution[1])\n",
    "    uv_points = torch.cartesian_prod(u_values, v_values) + 0.5 / torch.tensor(\n",
    "        uv_resolution,\n",
    "    )  # 0.5 is the center of a pixel\n",
    "    predictions = field(uv_points)\n",
    "    image_prediction = torch.clamp(predictions.reshape(uv_resolution[0], uv_resolution[1], 3), 0, 1)\n",
    "    image_prediction = image_prediction.permute(1, 0, 2)\n",
    "    rr.log(entity_path, rr.Image(image_prediction.numpy(force=True)))\n",
    "\n",
    "\n",
    "rr.set_time(\"iteration\", sequence=0)\n",
    "for i, field in enumerate(fields):\n",
    "    log_field_as_image(f\"field_{i}\", field, (-0.1, -0.1), (1.1, 1.1), (100, 100))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "## Train neural field\n",
    "\n",
    "Now we train the neural fields for a fixed number of iterations. If you run the cell twice, we continue training where we left off. To reset the fields, run the previous cell again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "field_ids = [0, 1, 2, 3]  # if you only want to train one of the fields\n",
    "num_iterations = 3000  # Run to 10_000 for better fit\n",
    "batch_size = 1000\n",
    "learning_rate = 1e-3\n",
    "log_image_period = 10\n",
    "\n",
    "response = requests.get(\"https://storage.googleapis.com/rerun-example-datasets/example_images/tiger.jpg\")\n",
    "# response = requests.get(\"https://storage.googleapis.com/rerun-example-datasets/example_images/bird.jpg\")\n",
    "src_array = np.asarray(PIL.Image.open(io.BytesIO(response.content)))\n",
    "target_image = torch.tensor(src_array).float() / 255\n",
    "\n",
    "rr.log(\"target\", rr.Image(target_image))\n",
    "\n",
    "try:\n",
    "    parameters = itertools.chain(*[fields[field_id].parameters() for field_id in field_ids])\n",
    "    optimizer = torch.optim.Adam(parameters, lr=learning_rate)\n",
    "    for iteration in tqdm(range(num_iterations)):\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        target_uvs = torch.rand(batch_size, 2)\n",
    "        target_jis = (target_uvs * torch.tensor([target_image.shape[1], target_image.shape[0]])).int()\n",
    "        target_rgbs = target_image[target_jis[:, 1], target_jis[:, 0]]\n",
    "        for field_id in field_ids:\n",
    "            field = fields[field_id]\n",
    "            total_iterations[field_id] += 1\n",
    "\n",
    "            predicted_rgbs = field(target_uvs)\n",
    "            loss = torch.nn.functional.mse_loss(target_rgbs, predicted_rgbs)\n",
    "\n",
    "            rr.set_time(\"iteration\", sequence=total_iterations[field_id])\n",
    "            rr.log(f\"loss/field_{field_id}\", rr.Scalars(loss.item()))\n",
    "            loss.backward()\n",
    "\n",
    "        optimizer.step()\n",
    "\n",
    "        if iteration % log_image_period == 0:\n",
    "            for field_id in field_ids:\n",
    "                log_field_as_image(f\"field_{field_id}\", fields[field_id], (-0.1, -0.1), (1.1, 1.1), (100, 100))\n",
    "except KeyboardInterrupt:\n",
    "    print(\"Training stopped.\")"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}
